package com.yifeng.repo.controller.rocketmq.message;

import com.gomcarter.frameworks.base.common.AssertUtils;
import com.gomcarter.frameworks.base.exception.CustomException;
import com.google.common.base.Strings;
import com.yifeng.repo.base.utils.common.BaseUtil;
import com.yifeng.repo.base.utils.converter.JacksonHelper;
import com.yifeng.repo.controller.rocketmq.configure.RocketmqProperties;
import com.yifeng.repo.controller.rocketmq.message.produce.OrderMessageQueueSelector;
import com.yifeng.repo.controller.rocketmq.message.produce.TargetMessageQueueSelector;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.acl.common.AclClientRPCHook;
import org.apache.rocketmq.acl.common.SessionCredentials;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.impl.producer.DefaultMQProducerImpl;
import org.apache.rocketmq.client.producer.*;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageQueue;

import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.*;

/**
 * Created by daibing on 2023/8/5.
 */
@Slf4j
public class ProduceManager {
    protected final ConcurrentMap<String, DefaultMQProducer> topic2Producer = new ConcurrentHashMap<>();
    protected final ConcurrentMap<String, TransactionMQProducer> topic2TxProducer = new ConcurrentHashMap<>();
    protected final ConcurrentMap<String, MessageQueueSelector> clazz2MessageQueueSelector = new ConcurrentHashMap<>();
    protected final String applicationName;
    protected final RocketmqProperties properties;

    public ProduceManager(String applicationName, RocketmqProperties properties) {
        this.applicationName = applicationName;
        this.properties = properties;
    }

    public void init() {
        log.info("yfcloud controller ProduceManager init success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    public void destroy() {
        topic2Producer.forEach((key, value) -> value.shutdown());
        topic2TxProducer.forEach((key, value) -> value.shutdown());
        log.info("yfcloud controller ProduceManager destroy success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    /**
     * 发送消息
     *
     * @param topic          消息主题
     * @param tag            消息标签用在 Consumer的场景中，用于服务端消息过滤。（补充：配置发送顺序消息后该方法默认使用tag作为路由键）
     * @param key            消息在业务层面的唯一标识，业务可以根据这个Key查找消息，要求Key唯一。注意RocketMQ不保证message id唯一
     * @param body           消息体
     * @param headers        消息头
     * @param delayTimeLevel 1 ~ 18，对应RocketMQ系统级别配置参数messageDelayLevel
     * @return msgId
     */
    public String sendMsg(String topic, String tag, String key, String body, Map<String, String> headers, Integer delayTimeLevel) {
        DefaultMQProducer producer = getProducer(topic);
        Message msg = new Message(topic, tag, key, body.getBytes(StandardCharsets.UTF_8));
        headers.forEach(msg::putUserProperty);
        if (delayTimeLevel != null) {
            msg.setDelayTimeLevel(delayTimeLevel);
        }
        try {
            SendResult result = producer.send(msg);
            AssertUtils.isTrue(result.getSendStatus() == SendStatus.SEND_OK, key + "发送失败: " + JacksonHelper.toJson(result));
            return result.getMsgId();
        } catch (Throwable t) {
            throw new CustomException(t);
        }
    }

    public String sendMsg(String topic, String tag, String key, String body) {
        return this.sendMsg(topic, tag, key, body, Collections.emptyMap(), null);
    }

    public String sendMsg(String topic, String tag, String key, String body, Map<String, String> headers) {
        return this.sendMsg(topic, tag, key, body, headers, null);
    }

    /**
     * 发送靶向消息
     *
     * @param topic          消息主题
     * @param tag            消息标签用在 Consumer的场景中，用于服务端消息过滤。（补充：配置发送顺序消息后该方法默认使用tag作为路由键）
     * @param key            消息在业务层面的唯一标识，业务可以根据这个Key查找消息，要求Key唯一。注意RocketMQ不保证message id唯一
     * @param body           消息体
     * @param headers        消息头
     * @param delayTimeLevel 1 ~ 18，对应RocketMQ系统级别配置参数messageDelayLevel
     * @param targetArg      靶向参数，决定选择消息队列，可以指定多个队列序号，逗号分隔，比如：1,3,7
     * @return msgId
     */
    public String sendTargetMsg(String topic, String tag, String key, String body, Map<String, String> headers,
                                Integer delayTimeLevel, Object targetArg) {
        DefaultMQProducer producer = getProducer(topic);
        Message msg = new Message(topic, tag, key, body.getBytes(StandardCharsets.UTF_8));
        headers.forEach(msg::putUserProperty);
        if (delayTimeLevel != null) {
            msg.setDelayTimeLevel(delayTimeLevel);
        }
        try {
            SendResult result = producer.send(msg, getMessageQueueSelector(TargetMessageQueueSelector.class), targetArg);
            AssertUtils.isTrue(result.getSendStatus() == SendStatus.SEND_OK, key + "发送失败: " + JacksonHelper.toJson(result));
            return result.getMsgId();
        } catch (Throwable t) {
            throw new CustomException(t);
        }
    }

    public String sendTargetMsg(String topic, String tag, String key, String body, Object targetArg) {
        return this.sendTargetMsg(topic, tag, key, body, Collections.emptyMap(), null, targetArg);
    }

    public String sendTargetMsg(String topic, String tag, String key, String body, Map<String, String> headers, Object targetArg) {
        return this.sendTargetMsg(topic, tag, key, body, headers, null, targetArg);
    }

    /**
     * 发送顺序消息
     *
     * @param topic          消息主题
     * @param tag            消息标签用在 Consumer的场景中，用于服务端消息过滤。（补充：配置发送顺序消息后该方法默认使用tag作为路由键）
     * @param key            消息在业务层面的唯一标识，业务可以根据这个Key查找消息，要求Key唯一。注意RocketMQ不保证message id唯一
     * @param body           消息体
     * @param headers        消息头
     * @param delayTimeLevel 1 ~ 18，对应RocketMQ系统级别配置参数messageDelayLevel
     * @param orderArg       顺序参数，决定选择消息队列，相同参数落到同一个队列就可以保证顺序，默认用tag
     * @return msgId
     */
    public String sendOrderMsg(String topic, String tag, String key, String body, Map<String, String> headers,
                               Integer delayTimeLevel, Object orderArg) {
        DefaultMQProducer producer = getProducer(topic);
        Message msg = new Message(topic, tag, key, body.getBytes(StandardCharsets.UTF_8));
        headers.forEach(msg::putUserProperty);
        if (delayTimeLevel != null) {
            msg.setDelayTimeLevel(delayTimeLevel);
        }
        try {
            SendResult result = producer.send(msg, getMessageQueueSelector(OrderMessageQueueSelector.class), orderArg);
            AssertUtils.isTrue(result.getSendStatus() == SendStatus.SEND_OK, key + "发送失败: " + JacksonHelper.toJson(result));
            return result.getMsgId();
        } catch (Throwable t) {
            throw new CustomException(t);
        }
    }

    public String sendOrderMsg(String topic, String tag, String key, String body, Object orderArg) {
        return this.sendOrderMsg(topic, tag, key, body, Collections.emptyMap(), null, orderArg);
    }

    public String sendOrderMsg(String topic, String tag, String key, String body, Map<String, String> headers, Object orderArg) {
        return this.sendOrderMsg(topic, tag, key, body, headers, null, orderArg);
    }

    /**
     * 发送自定义消息
     *
     * @param topic          消息主题
     * @param tag            消息标签用在 Consumer的场景中，用于服务端消息过滤。（补充：配置发送顺序消息后该方法默认使用tag作为路由键）
     * @param key            消息在业务层面的唯一标识，业务可以根据这个Key查找消息，要求Key唯一。注意RocketMQ不保证message id唯一
     * @param body           消息体
     * @param headers        消息头
     * @param delayTimeLevel 1 ~ 18，对应RocketMQ系统级别配置参数messageDelayLevel
     * @param customArg      自定义参数
     * @return msgId
     */
    public String sendCustomMsg(String topic, String tag, String key, String body, Map<String, String> headers,
                                Integer delayTimeLevel, Object customArg) {
        DefaultMQProducer producer = getProducer(topic);
        Message msg = new Message(topic, tag, key, body.getBytes(StandardCharsets.UTF_8));
        headers.forEach(msg::putUserProperty);
        if (delayTimeLevel != null) {
            msg.setDelayTimeLevel(delayTimeLevel);
        }
        try {
            SendResult result = producer.send(msg, getMessageQueueSelector(properties.getCustomMessageQueueSelectorClass()), customArg);
            AssertUtils.isTrue(result.getSendStatus() == SendStatus.SEND_OK, key + "发送失败: " + JacksonHelper.toJson(result));
            return result.getMsgId();
        } catch (Throwable t) {
            throw new CustomException(t);
        }
    }

    public String sendCustomMsg(String topic, String tag, String key, String body, Object customArg) {
        return this.sendCustomMsg(topic, tag, key, body, Collections.emptyMap(), null, customArg);
    }

    public String sendCustomMsg(String topic, String tag, String key, String body, Map<String, String> headers, Object customArg) {
        return this.sendCustomMsg(topic, tag, key, body, headers, null, customArg);
    }

    /**
     * 发送事务消息（不支持延时消息和批量消息）
     *
     * @param topic         消息主题
     * @param tag           消息标签用在 Consumer的场景中，用于服务端消息过滤。（补充：配置发送顺序消息后该方法默认使用tag作为路由键）
     * @param key           消息在业务层面的唯一标识，业务可以根据这个Key查找消息，要求Key唯一。注意RocketMQ不保证message id唯一
     * @param body          消息体
     * @param headers       消息头
     * @param transactionId 事务id
     * @param txArg         事务参数
     * @return msgId
     */
    public String sendTxMsg(String topic, String tag, String key, String body, Map<String, String> headers,
                            String transactionId, Object txArg) {
        TransactionMQProducer producer = getTxProducer(topic);
        Message msg = new Message(topic, tag, key, body.getBytes(StandardCharsets.UTF_8));
        headers.forEach(msg::putUserProperty);
        if (!Strings.isNullOrEmpty(transactionId)) {
            msg.setTransactionId(transactionId);
        }
        try {
            SendResult result = producer.sendMessageInTransaction(msg, txArg);
            AssertUtils.isTrue(result.getSendStatus() == SendStatus.SEND_OK, key + "发送失败: " + JacksonHelper.toJson(result));
            return result.getMsgId();
        } catch (Throwable t) {
            throw new CustomException(t);
        }
    }

    public String sendTxMsg(String topic, String tag, String key, String body, String transactionId, Object customArg) {
        return this.sendTxMsg(topic, tag, key, body, Collections.emptyMap(), transactionId, customArg);
    }

    /**
     * 发送延迟消息：
     * 1、RocketMQ将延时队列的延时延时时间分为18个级别
     *    1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 分别对应下面的延迟时间，在使用时，直接传递 level即可。
     * 2、RocketMQ系统级别配置，默认如下，支持修改
     *    messageDelayLevel=1s 5s 10s 30s 1m 2m 3m 4m 5m 6m 7m 8m 9m 10m 20m 30m 1h 2h
     */

    public DefaultMQProducer getDefaultMQProducer(String topic) {
        return getProducer(topic);
    }

    public TransactionMQProducer getTransactionMQProducer(String topic) {
        return getTxProducer(topic);
    }

    /**
     * 查询指定topic的消息队列
     */
    public List<MessageQueue> listMessageQueue(String topic) {
        DefaultMQProducer defaultMQProducer = topic2Producer.get(topic);
        if (defaultMQProducer == null) {
            defaultMQProducer = topic2TxProducer.get(topic);
        }
        AssertUtils.notNull(defaultMQProducer, "获取不到Topic对应的生产者：" + topic);
        try {
            return defaultMQProducer.fetchPublishMessageQueues(topic);
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
    }

    /**
     * 获取队列的统计数据
     *
     * @param topic     消息主题
     * @param timestamp 开始时间，单位毫秒
     * @return queueId -> searchOffset
     */
    public Map<Integer, Long> listMessageQueueStats(String topic, long timestamp) {
        DefaultMQProducer defaultMQProducer = topic2Producer.get(topic);
        if (defaultMQProducer == null) {
            defaultMQProducer = topic2TxProducer.get(topic);
        }
        AssertUtils.notNull(defaultMQProducer, "获取不到Topic对应的生产者：" + topic);
        Map<Integer, Long> mqStats = new HashMap<>();
        try {
            List<MessageQueue> messageQueues = defaultMQProducer.fetchPublishMessageQueues(topic);
            for (MessageQueue messageQueue : messageQueues) {
                mqStats.put(messageQueue.getQueueId(), defaultMQProducer.searchOffset(messageQueue, timestamp));
            }
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
        return mqStats;
    }

    protected DefaultMQProducer getProducer(String topic) {
        if (properties.isMasterTopicOnlyProduce() && !properties.getMasterTopic().equals(topic)) {
            throw new CustomException("当前仅支持发布消息到应用自己的topic：" + topic);
        }
        DefaultMQProducer producer = topic2Producer.get(topic);
        if (producer != null) {
            return producer;
        }
        synchronized (topic2Producer) {
            producer = topic2Producer.get(topic);
            if (producer != null) {
                return producer;
            }
            producer = buildNewProducer(topic);
            topic2Producer.put(topic, producer);
            return producer;
        }
    }

    protected TransactionMQProducer getTxProducer(String topic) {
        if (properties.isMasterTopicOnlyProduce() && !properties.getMasterTopic().equals(topic)) {
            throw new CustomException("当前仅支持发布消息到应用自己的topic：" + topic);
        }
        TransactionMQProducer txProducer = topic2TxProducer.get(topic);
        if (txProducer != null) {
            return txProducer;
        }
        synchronized (topic2TxProducer) {
            txProducer = topic2TxProducer.get(topic);
            if (txProducer != null) {
                return txProducer;
            }
            txProducer = buildNewTxProducer(topic);
            topic2TxProducer.put(topic, txProducer);
            return txProducer;
        }
    }

    protected DefaultMQProducer buildNewProducer(String topic) {
        RocketmqProperties.TopicProperties topicProperties = properties.getTopicPropertiesList()
                .stream()
                .filter(s -> s.getTopic().equals(topic))
                .findFirst()
                .orElseThrow(() -> new CustomException("没有Topic配置：" + topic));
        AclClientRPCHook rpcHook = new AclClientRPCHook(new SessionCredentials(topicProperties.getAccessKey(), topicProperties.getSecretKey()));
        DefaultMQProducer producer = new DefaultMQProducer(topicProperties.getProducerGroup(), rpcHook);
        producer.setNamesrvAddr(topicProperties.getNamesrvAddr());
        producer.setSendMsgTimeout(topicProperties.getSendMsgTimeout());
        producer.setRetryTimesWhenSendAsyncFailed(topicProperties.getRetryTimesWhenSendAsyncFailed());
        producer.setRetryTimesWhenSendFailed(topicProperties.getRetryTimesWhenSendFailed());
        producer.setRetryAnotherBrokerWhenNotStoreOK(topicProperties.isRetryAnotherBrokerWhenNotStoreOK());
        producer.setMaxMessageSize(topicProperties.getMaxMessageSize());
        producer.setCompressMsgBodyOverHowmuch(topicProperties.getCompressMsgBodyOverHowmuch());
        try {
            producer.start();
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
        log.info("buildNewProducer success：{}", topic);
        return producer;
    }

    /**
     * @see DefaultMQProducerImpl#initTransactionEnv
     */
    protected TransactionMQProducer buildNewTxProducer(String topic) {
        RocketmqProperties.TopicProperties topicProperties = properties.getTopicPropertiesList()
                .stream()
                .filter(s -> s.getTopic().equals(topic))
                .findFirst()
                .orElseThrow(() -> new CustomException("没有Topic配置：" + topic));
        AssertUtils.notNull(properties.getTransaction().getTransactionListenerClass(), "未设置事务消息监听器！");

        AclClientRPCHook rpcHook = new AclClientRPCHook(new SessionCredentials(topicProperties.getAccessKey(), topicProperties.getSecretKey()));
        TransactionMQProducer producer = new TransactionMQProducer(topicProperties.getProducerGroup(), rpcHook);
        producer.setNamesrvAddr(topicProperties.getNamesrvAddr());
        producer.setSendMsgTimeout(topicProperties.getSendMsgTimeout());
        producer.setRetryTimesWhenSendFailed(topicProperties.getRetryTimesWhenSendFailed());
        producer.setRetryTimesWhenSendAsyncFailed(topicProperties.getRetryTimesWhenSendAsyncFailed());
        producer.setRetryAnotherBrokerWhenNotStoreOK(topicProperties.isRetryAnotherBrokerWhenNotStoreOK());
        producer.setMaxMessageSize(topicProperties.getMaxMessageSize());
        producer.setCompressMsgBodyOverHowmuch(topicProperties.getCompressMsgBodyOverHowmuch());
        producer.setExecutorService(new ThreadPoolExecutor(
                properties.getTransaction().getCheckThreadPoolMinSize(),
                properties.getTransaction().getCheckThreadPoolMaxSize(),
                1000 * 60,
                TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<>(properties.getTransaction().getCheckRequestHoldMax()),
                new ThreadFactory() {
                    final Random random = new Random();
                    @Override
                    public Thread newThread(Runnable r) {
                        Thread thread = new Thread(r);
                        thread.setName(topic + "-tx-" + this.random.nextInt());
                        return thread;
                    }
                },
                new ThreadPoolExecutor.AbortPolicy())
        );
        try {
            producer.setTransactionListener(properties.getTransaction().getTransactionListenerClass().newInstance());
            producer.start();
        } catch (Exception e) {
            throw new CustomException(e);
        }
        log.info("buildNewTxProducer success: {}", topic);
        return producer;
    }

    protected MessageQueueSelector getMessageQueueSelector(Class<? extends MessageQueueSelector> messageQueueSelectorClass) {
        MessageQueueSelector messageQueueSelector = clazz2MessageQueueSelector.get(messageQueueSelectorClass.getName());
        if (messageQueueSelector != null) {
            return messageQueueSelector;
        }
        synchronized (clazz2MessageQueueSelector) {
            messageQueueSelector = clazz2MessageQueueSelector.get(messageQueueSelectorClass.getName());
            if (messageQueueSelector != null) {
                return messageQueueSelector;
            }
            try {
                messageQueueSelector = messageQueueSelectorClass.newInstance();
            } catch (Exception e) {
                throw new CustomException(e);
            }
            clazz2MessageQueueSelector.put(messageQueueSelectorClass.getName(), messageQueueSelector);
            return messageQueueSelector;
        }
    }

}
